/*
 * Alloy Analyzer
 * Copyright (c) 2002 Massachusetts Institute of Technology
 *
 * The Alloy Analyzer is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * The Alloy Analyzer is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with the Alloy Analyzer; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

package alloy.transform;

import java.util.Collection;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;

import alloy.ast.ArrowMultExpr;
import alloy.ast.BinaryExpr;
import alloy.ast.Decl;
import alloy.ast.Expr;
import alloy.ast.Formula;
import alloy.ast.FormulaSeq;
import alloy.ast.Formulas;
import alloy.ast.Id;
import alloy.ast.LeafId;
import alloy.ast.Multiplicity;
import alloy.ast.MultiplicityExpr;
import alloy.ast.MultiplicityFormula;
import alloy.ast.QuantifiedExpr;
import alloy.ast.QuantifiedFormula;
import alloy.ast.SetMultExpr;
import alloy.ast.SigExpr;
import alloy.ast.Signature;
import alloy.ast.Variable;
import alloy.ast.VariableExpr;
import alloy.util.Dbg;

/**
 * <p>Contains methods to generate constraints implied by certain Alloy
 * constructs.  These constructs include <code>extends</code>, <code>in</code>, 
 * multiplicity markings in declarations, set membership of declarations, and 
 * qualifiers such as abstract.
 */
public class ImplicitConstraintGenerator {

    /** a mapping from each Signature to the list of implicit constraints for that Signature. */
    public Map sigToConstraints = new LinkedHashMap(); // String -> Formulas

    /** a mapping from the {@link alloy.ast.LeafId}s of fields to the {@link MultiplicityExpr} of
        the field declaration */
    public Map leafIdToMult = new LinkedHashMap(); // LeafId -> MultiplicityExpr
    
    private Signature _univ;

    /**
     * Generate formulas expressing the 
     * implicit constraints for each Signature and add them to the mapping.
     * @param sigs_ a collection of all signatures found in the spec
     * being analyzed
     */
    public void createSigConstraints(Collection sigs_) {
	
	_univ = _getUniv(sigs_);
	
	for (Iterator sigIter = sigs_.iterator(); sigIter.hasNext();) {
            
            Signature sig = (Signature) sigIter.next();
	    Formulas sigConstraints = new Formulas();

	    // generate LeafId->MultiplicityExpr mappings for sig and its fields,
	    // and add them to leafIdToMult
	    _updateLeafIdToMult(sig);

	    // generate constraints implied by field declarations in sig and add them
	    // to sigConstraints
	    _addFieldDeclConstraints(sig, sigConstraints);

	    // add multiplicity constraints for sig to sigConstraints
			if (!sig.getMult().equals(Multiplicity.SET)) {
				sigConstraints.addFormula(DesugarMultFormulasVisitor.genSetMultConstraint(sig.getMult(), sig.toSigExpr()));
			}

	    // add extension or subset constraints for sig to sigConstraints
	    _addExtensionConstraints(sig, sigConstraints);

	    sigToConstraints.put(sig, sigConstraints);

            //System.out.println("In createSigConstraints: "  + sigName + " constraints " + sigConstraints.nodeString());
        }

    }

    /**
     * @return univ
     **/
	private static Signature _getUniv(Collection sigs) {
		for (Iterator sigIter = sigs.iterator(); sigIter.hasNext();) {
			Signature sig = (Signature) sigIter.next();
			if (!sig.isSubset() && sig.getSigType().getSuperType() == null)
				return sig;
		}
		Dbg.fatal("univ undefined");
		return null;
	}

    /** 
     * Generates LeafId->MultiplicityExpr mappings for sig_ and all of its fields,
     * and adds them to leafIdToMult. 
     **/
    private void _updateLeafIdToMult(Signature sig_) {
	Dbg.check(sig_.getLeafId() == null); // leaf id should not be set at this point
	
	// generate leaf id for sig_
	sig_.setLeafId(new LeafId());

	// generate leaf ids for the fields of the signature
	sig_.setLeafIdMap(sig_.getLocalScope().generateLeafIdMapping(null));

	// add multiplicity expr for sig_ to leafIdToMult
	leafIdToMult.put(sig_.getLeafId(), 
			 new SetMultExpr((Multiplicity) sig_.getMult().copy(), sig_.toSigExpr()));

	// add multiplicity exprs for sig_'s fields to leafIdToMult
        Iterator declIter = sig_.getDecls().getDeclIter();
        while (declIter.hasNext()) {
            Decl decl = (Decl)declIter.next();
            MultiplicityExpr multExpr = decl.getExpr();
            Iterator varIter = decl.getVariables().getVariableIter();
            while (varIter.hasNext()) {
                LeafId id = sig_.getFieldLeafId(((Variable)varIter.next()).nodeString());
                // NOTE: not sure if this check is needed
                if (!leafIdToMult.containsKey(id)) {
                    leafIdToMult.put(id, createDesugaredMultExpr(sig_.toSigExpr(), multExpr));
                }
            }
        }
    }

    /** 
     * Generate constraints implied by field declarations in sig_ and add them
     * to sigConstraints_ 
     **/
    private void _addFieldDeclConstraints(Signature sig_, Formulas sigConstraints_) {
	
        Formulas fieldConstraints = new Formulas();
	
	// create "this" variable
		VariableExpr sigVarExpr = new VariableExpr(new Variable(new Id("this")));
		sigVarExpr.setType(sig_.getUnionType());

		// generate constraints for declarations
		Iterator declIter = sig_.getDecls().getDeclIter();
		while (declIter.hasNext()) {
			Decl decl = (Decl) declIter.next();
			_handleFieldDecl(sigConstraints_, fieldConstraints, sig_.toSigExpr(), sigVarExpr, decl);
		}

	// add constraints for declarations, if any
        if (!fieldConstraints.isLeaf()) { // some constraints to be added
	    fieldConstraints.setLocation(sig_.getDecls().getLocation());

            Variable var = new Variable(new Id("this"));
            SetMultExpr multExpr = new SetMultExpr(Multiplicity.NONE, sig_.toSigExpr());
            multExpr.setType(sig_.getUnionType());
            Decl decl = Decl.in(var, multExpr);
            QuantifiedFormula quantForm = QuantifiedFormula.all(decl, new FormulaSeq(fieldConstraints));
            quantForm.annotate("constraints for fields of signature " + sig_.getName().nodeString());
            quantForm.setLocation(sig_.getDecls().getLocation());
            sigConstraints_.addFormula(quantForm);

        }

    }
	
    private void _handleFieldDecl(
    		Formulas sigConstraints,
    		Formulas fieldConstraints, 
		SigExpr sigExpr,
		VariableExpr sigVarExpr, Decl decl) {

		Iterator fields = decl.getVariables().getVariableIter();
		MultiplicityExpr multExpr = decl.getExpr();
		Set fieldExprs = new LinkedHashSet();

		while (fields.hasNext()) {
			Variable field = (Variable) fields.next();

			// variable expr for field, used in constraints
			VariableExpr fieldVarExpr = new VariableExpr(field);
			fieldVarExpr.setType(sigVarExpr.getType().product(multExpr.getType()));
			fieldVarExpr.setLocation(field.getLocation());

			// this.f, used in constraints
			BinaryExpr sigFieldExpr = ((Expr) sigVarExpr.copy()).join(fieldVarExpr);
			sigFieldExpr.setLocation(field.getLocation());

			fieldExprs.add(sigFieldExpr);

			MultiplicityExpr multExprCopy = (MultiplicityExpr) multExpr.copy();

			// in a field decl, set multiplicity of NONE means ONE
			if (multExprCopy.getType().arity() == 1
					&& multExprCopy.getMult().equals(Multiplicity.NONE))
				multExprCopy.setMult(Multiplicity.ONE);

			// add multiplicity formula for f ===> this.f : multExpr
			Formula multFormula =
				new MultiplicityFormula((Expr) sigFieldExpr.copy(), multExprCopy);

			multFormula.annotate("constraints for field " + field.nodeString());
			multFormula.setLocation(field.getLocation());
			fieldConstraints.addFormula(multFormula);

			// add membership constraint for f
			sigConstraints.addFormula(_genFieldMembershipConstraint(sigExpr, multExpr, fieldVarExpr));
		}
		// generate necessary constraints for qualifiers
		if (decl.isDisj() && fieldExprs.size() > 1) { //disjoint and more than one field declared
			fieldConstraints.addFormula(genDisjConstraints(fieldExprs));
		}
		if (decl.isExh() && fieldExprs.size() > 0) {
			fieldConstraints.addFormula(genExhConstraints(fieldExprs, multExpr.toExpr()));
		}
    }
    
    /**
	 * generate membership constraint for f. For example, consider the
	 * declaration sig A { f1 : expr }
	 * 
	 * 
	 * for the above sig, we would generate f1 in A->type(expr)
	 * 
	 * @param sigExpr_
	 *            SigExpr that represents the signature containing the field (A)
	 * @param multExpr_
	 *            MultiplicityExpr that represents the expression defining the
	 *            field (expr)
	 * @param fieldExpr_
	 *            VariableExpr that represents the field (f1)
	 */
    private Formula _genFieldMembershipConstraint(SigExpr sigExpr_, MultiplicityExpr multExpr_,
						  VariableExpr fieldExpr_) {
	
	// type(expr)
	Expr multTypeExpr = multExpr_.getType().toExpr();
	
	// A->type(expr)
	BinaryExpr productExpr = ((SigExpr)sigExpr_.copy()).product(multTypeExpr);

	// f1 in A->type(expr)
	Formula constraint = ((Expr) fieldExpr_.copy()).in(productExpr);
	constraint.annotate("field membership constraint " + fieldExpr_.nodeString());
        constraint.setLocation(fieldExpr_.getLocation());
        return constraint;

    }

	/**
	 * Add extension constraints for sig_ to sigConstraints_
	 */
	private void _addExtensionConstraints(Signature sig_,
			Formulas sigConstraints_) {

		if (sig_.isSubset()) {
			// if sig_ is a subset, add the constraint sig_ in superset1 +
			// supserset2 + . . .
			Expr sumOfParents = _genUnionExpr(sig_.getSubsets().getSigExprIter());
			sigConstraints_.addFormula(_genSigSubsetConstraint(sig_, sumOfParents));

		} else {
			// if sig_ extends another sig, add the constraint sig_ in parent
			if (!sig_.getExtends().isLeaf()) {
				SigExpr parent = (SigExpr) ((SigExpr) sig_.getExtends()	.getSigExprIter().next()).copy();
				sigConstraints_.addFormula(_genSigSubsetConstraint(sig_, parent));
			}

			// if sig_ is extended by other sigs ...
			if (sig_.getExtByIter().hasNext()) {
				// create a set of SigExprs corresponding to Signatures that
				// extend sig_
				Set extenders = new LinkedHashSet();
				for (Iterator extendedBy = sig_.getExtByIter(); extendedBy.hasNext();) {
					Signature child = (Signature) extendedBy.next();
					extenders.add(child.toSigExpr());
				}

				// if sig_ is extended by more than one other sig, then
				// add the constraint that all signatures extending sig_ are
				// disjoint
				if (extenders.size() > 1) {
					sigConstraints_.addFormula(genDisjConstraints(extenders));
				}

				// if sig_ is abstract then
				// add the constraint that the signature(s) extending sig_ also
				// partition it,
				if (sig_.isAbstract()) {
					Expr sumOfChildren = _genUnionExpr(extenders.iterator());
					sigConstraints_.addFormula(_genSigSubsetConstraint(sig_, sumOfChildren));
				}
			}
		}
	}

    /**
     * @return a Formula which states that sig_ is a subset of expr_:  sig_ in expr_
     **/
	private static Formula _genSigSubsetConstraint(Signature sig_, Expr expr_) {
		Formula subsetFormula = sig_.toSigExpr().in(expr_);
		subsetFormula.annotate("subset constraint for " + sig_.getName().nodeString());
		subsetFormula.setLocation(sig_.getLocation());
		return subsetFormula;
	}

    /** 
     * generates a union expression from copies of expressions returned 
     * by exprIter_ 
     **/
    private static Expr _genUnionExpr(Iterator exprIter_) {
        Dbg.chk(exprIter_.hasNext());
        Expr retExpr = (Expr) ((Expr) exprIter_.next()).copy();
	while (exprIter_.hasNext()) {
            retExpr = retExpr.union((Expr) ((Expr)exprIter_.next()).copy());
        }
        return retExpr;
    }


    /**
     * Generate a formula expressing the constraint that 
     * a set of expressions is disjoint.
     */
    public static Formula genDisjConstraints(Set exprs_) {
        Formulas constraints = new Formulas();
        Set tempSet = new LinkedHashSet(exprs_);
        Iterator exprIter = exprs_.iterator();
        while (exprIter.hasNext()) {
            Expr expr = (Expr) exprIter.next();
            tempSet.remove(expr);
            BinaryExpr binExpr = ((Expr)expr.copy()).intersect(_genUnionExpr(tempSet.iterator()));
            constraints.addFormula(QuantifiedExpr.no(binExpr));
            tempSet.add(expr);
        }
        Formula ret = new FormulaSeq(constraints);
        ret.annotate("disj/part constraint");
        return ret;
    }

    /**
     * Generate a formula expressing the constraint that a 
     * set of expressions are exhaustive of another expression.
     */
    public static Formula genExhConstraints(Set exprs_, Expr multExpr_) {
        Formula ret = multExpr_.eq(_genUnionExpr(exprs_.iterator()));
        ret.annotate("part/exh constraint");
        return ret;
    }

	
    /**
     * create a multiplicity expr which reflects the prefix expr followed
     * by the expression in an existing multiplicity expr.  preserves
     * the multiplicities in the existing multiplicity expr
     */
	public static MultiplicityExpr createDesugaredMultExpr(Expr prefixExpr_, 	MultiplicityExpr multExpr_) {
		MultiplicityExpr newMultExpr;
		if (prefixExpr_ != null) {
			newMultExpr = new ArrowMultExpr(
					Multiplicity.NONE,
					new SetMultExpr(Multiplicity.NONE, (Expr) prefixExpr_.copy()),
					(MultiplicityExpr) multExpr_.copy());
			newMultExpr.setType(prefixExpr_.getType().product(multExpr_.getType()));
		} else {
			newMultExpr = (MultiplicityExpr) multExpr_.copy();
		}
		return newMultExpr;
	}
    
}
